module Timing

export @timed_function, TimingStats, print_timing_report, reset_timing!, get_timing_stats

using Printf

# Global timing storage
const timing_data = Dict{String, Vector{Float64}}()
const timing_hierarchy = Dict{String, Vector{String}}()
const call_counts = Dict{String, Int}()
const current_stack = String[]

mutable struct TimingStats
    name::String
    calls::Int
    total_time::Float64
    max_time::Float64
    min_time::Float64
    mean_time::Float64
end

"""
    @timed_function name expr

Times the execution of expr and stores statistics under name.
Similar to GPTL's gptlstart/gptlstop.
"""
macro timed_function(name, expr)
    quote
        local start_time = time()
        push!(current_stack, $(esc(name)))

        # Track hierarchy
        if length(current_stack) > 1
            parent = current_stack[end-1]
            if !haskey(timing_hierarchy, parent)
                timing_hierarchy[parent] = String[]
            end
            if !($(esc(name)) in timing_hierarchy[parent])
                push!(timing_hierarchy[parent], $(esc(name)))
            end
        end

        local result = $(esc(expr))

        local elapsed = time() - start_time

        # Store timing data
        if !haskey(timing_data, $(esc(name)))
            timing_data[$(esc(name))] = Float64[]
            call_counts[$(esc(name))] = 0
        end
        push!(timing_data[$(esc(name))], elapsed)
        call_counts[$(esc(name))] += 1

        pop!(current_stack)

        result
    end
end

"""
    reset_timing!()

Clear all timing data.
"""
function reset_timing!()
    empty!(timing_data)
    empty!(timing_hierarchy)
    empty!(call_counts)
    empty!(current_stack)
end

"""
    get_timing_stats(name::String)

Get timing statistics for a specific function.
"""
function get_timing_stats(name::String)
    if !haskey(timing_data, name)
        return nothing
    end

    times = timing_data[name]
    return TimingStats(
        name,
        call_counts[name],
        sum(times),
        maximum(times),
        minimum(times),
        sum(times) / length(times)
    )
end

"""
    print_timing_report(; indent_level=0, parent="main", show_all=false)

Print a hierarchical timing report similar to GPTL output.
"""
function print_timing_report(; indent_level::Int=0, parent::String="main", show_all::Bool=false)
    if indent_level == 0
        println("\n" * "="^80)
        println("Julia FLEXINVERT Timing Report (GPTL-style)")
        println("="^80)
        println("                         Called  Recurse Wallclock max       min       ")
    end

    # Print current level
    indent = "  " * repeat("  ", indent_level)

    if haskey(timing_data, parent)
        stats = get_timing_stats(parent)
        if stats !== nothing
            # Format output similar to GPTL
            if stats.total_time >= 0.01 || show_all
                @printf("%s%-24s %5d    -    %8.3f  %8.3f  %8.3f\n",
                    indent, parent, stats.calls, stats.total_time,
                    stats.max_time, stats.min_time)
            else
                @printf("%s%-24s %5d    -    %8.2e  %8.2e  %8.2e\n",
                    indent, parent, stats.calls, stats.total_time,
                    stats.max_time, stats.min_time)
            end

            # Print children
            if haskey(timing_hierarchy, parent)
                for child in timing_hierarchy[parent]
                    print_timing_report(indent_level=indent_level+1, parent=child, show_all=show_all)
                end
            end
        end
    end

    if indent_level == 0
        println("="^80)

        # Print summary statistics
        total_time = sum(sum(v) for v in values(timing_data))
        total_calls = sum(values(call_counts))

        println("\nSummary:")
        @printf("  Total wallclock time: %.3f seconds\n", total_time)
        @printf("  Total function calls: %d\n", total_calls)

        # Memory usage estimate (Julia-specific)
        mem_usage = Base.summarysize(timing_data) / 1024 / 1024
        @printf("  Timing overhead memory: %.2f MB\n", mem_usage)
    end
end

# Export convenience function for timing blocks
"""
    with_timing(name::String, f::Function)

Execute function f while timing it under the given name.
"""
function with_timing(name::String, f::Function)
    @timed_function name f()
end

end # module